package com.embeddings.controller;


import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.embeddings.bean.SearchResult;
import com.embeddings.embed.EmbedClient;
import com.embeddings.embed.ReRankClient;
import com.embeddings.qwen.Main;
import com.embeddings.service.EsDocumentService;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/embedding")
public class DataProcessController {

    @Value("${qwen.api-key}")
    private String apiKey;

    @Value("${qwen.model}")
    private String model;

    @Value("${embedding.uri}")
    private String embeddingUri;

    @Value("${embedding.api-key}")
    private String embeddingApiKey;

    @Value("${re-rank.uri}")
    private String ReRankUri;

    @Value("${re-rank.api-key}")
    private String ReRankApiKey;

    @Resource
    private EsDocumentService service;

    @PostMapping("/process")
    public List<SearchResult> process(@RequestBody Map<String, String> dto) throws IOException {
        String keyword = dto.get("keyword");
        return service.searchWithGivenVector(EmbedClient.getEmbedding(embeddingUri, embeddingApiKey, keyword));
    }

    @PostMapping("/generate")
    public String generate(@RequestBody Map<String, String> dto) throws IOException, NoApiKeyException, InputRequiredException {
        String keyword = dto.get("keyword");
        List<SearchResult> searchResults = service.searchWithGivenVector(EmbedClient.getEmbedding(embeddingUri, embeddingApiKey, keyword));
        List<String> collect = searchResults.stream().map(searchResult -> searchResult.getSell().getRemark()).toList();
        System.out.println("搜索结果searchResults = " + collect);

        // 重排处理
        List<String> contentList = new ArrayList<>();
        searchResults.forEach(searchResult -> contentList.add(searchResult.getSell().getRemark()));
        String reRank = ReRankClient.reRank(ReRankUri, ReRankApiKey, contentList, keyword);
        System.out.println("重排结果reRank = " + reRank);

        JSONObject jsonObject = JSON.parseObject(reRank, JSONObject.class);
        Object reRankPassages = jsonObject.get("rerank_passages");

        //增强生成
        GenerationResult result = Main.callWithMessage(model, apiKey, keyword, reRankPassages.toString());
        System.err.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~回答开始~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
        System.out.println("千问：" + result.getOutput().getChoices().get(0).getMessage().getContent());
        return result.getOutput().getChoices().get(0).getMessage().getContent();
    }
}
